import torch
import torch.nn as nn
import torch.nn.functional as F


class ResNet(nn.Module):
    def __init__(self, block, layers):
        super().__init__()
        self.in_channels = 64

        # 初始卷积层
        self.conv1 = nn.Conv2d(
            3, 64, kernel_size=7,
            stride=2, padding=3, bias=False
        )
        self.bn1 = nn.BatchNorm2d(64)
        self.relu = nn.ReLU(inplace=True)
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)

        # 构建四个残差阶段
        self.layer1 = self._make_layer(block, 64, layers[0], stride=1)
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)



        # 初始化参数
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def _make_layer(self, block, out_channels, num_blocks, stride):
        layers = []
        # 每个阶段的第一个块处理下采样
        layers.append(block(self.in_channels, out_channels, stride))
        self.in_channels = out_channels * block.expansion
        # 后续块保持通道数不变
        for _ in range(1, num_blocks):
            layers.append(block(self.in_channels, out_channels))
        return nn.Sequential(*layers)

    def forward(self, x):
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x)
        x = self.layer1(x)
        tezheng1 = self.layer2(x)
        tezheng2 = self.layer3(tezheng1)
        tezheng3 = self.layer4(tezheng2)
        return tezheng1,tezheng2,tezheng3




class BottleneckBlock(nn.Module):
    expansion = 2 # 输出通道扩展系数

    def __init__(self, in_channels, out_channels, stride=1):
        super().__init__()
        mid_channels = out_channels // self.expansion

        self.conv1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(mid_channels)
        self.conv2 = nn.Conv2d(
            mid_channels, mid_channels, kernel_size=3,
            stride=stride, padding=1, bias=False
        )
        self.bn2 = nn.BatchNorm2d(mid_channels)
        self.conv3 = nn.Conv2d(
            mid_channels, out_channels * self.expansion,
            kernel_size=1, bias=False
        )
        self.bn3 = nn.BatchNorm2d(out_channels * self.expansion)
        self.relu = nn.ReLU(inplace=True)

        # 跳跃连接处理维度变化
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels * self.expansion:
            self.shortcut = nn.Sequential(
                nn.Conv2d(
                    in_channels, out_channels * self.expansion,
                    kernel_size=1, stride=stride, bias=False
                ),
                nn.BatchNorm2d(out_channels * self.expansion)
            )

    def forward(self, x):
        identity = x
        x = self.relu(self.bn1(self.conv1(x)))
        x = self.relu(self.bn2(self.conv2(x)))
        x = self.bn3(self.conv3(x))
        x += self.shortcut(identity)
        x = self.relu(x)
        return x

def ResNet50(num_classes=1000):
    return ResNet(BottleneckBlock, [3, 4, 6, 3])


import torch.nn as nn


class FCBlock(nn.Module):
    """
    全连接模块封装类
    特征：线性层 + BatchNorm1d + GELU激活的三件套组合（最后一层不含激活）

    参数说明：
        input_dim (int): 输入特征维度
        hidden_dims (list): 隐藏层维度列表，如[128,256]表示两个隐藏层
        output_dim (int): 最终输出维度
    """

    def __init__(self, input_dim=42, hidden_dims=None, output_dim=512):
        super().__init__()
        if hidden_dims is None:
            hidden_dims = [128, 256]
        layers = []
        in_features = input_dim

        # 动态构建隐藏层
        for h_dim in hidden_dims:
            layers.extend([
                nn.Linear(in_features, h_dim),
                nn.BatchNorm1d(h_dim),
                nn.GELU()
            ])
            in_features = h_dim  # 更新输入维度

        # 添加输出层（不含激活函数）
        layers.extend([
            nn.Linear(in_features, output_dim),
            nn.BatchNorm1d(output_dim)
        ])

        self.block = nn.Sequential(*layers)

        # 初始化参数
        self._init_weights()

    def _init_weights(self):
        """ He初始化线性层权重，BN层保持默认 """
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='gelu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        return self.block(x)

class ChannelAttention(nn.Module):
    def __init__(self, in_planes, ratio=16):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)

        self.fc1 = nn.Conv2d(in_planes, in_planes // ratio, 1, bias=False)
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Conv2d(in_planes // ratio, in_planes, 1, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = self.fc2(self.relu1(self.fc1(self.avg_pool(x))))
        max_out = self.fc2(self.relu1(self.fc1(self.max_pool(x))))
        out = avg_out + max_out
        channel_weights = self.sigmoid(out)
        return x * channel_weights.expand_as(x)



class Conv(nn.Module):
    """深度可分离卷积模块 (参考网页4的GSConvE方案)"""

    def __init__(self, in_ch, out_ch, kernel_size=3, depthwise=False):
        super().__init__()
        groups = in_ch if depthwise else 1
        self.conv = nn.Sequential(
            nn.Conv2d(in_ch, in_ch, kernel_size,
                      padding=kernel_size // 2, groups=groups),
            nn.BatchNorm2d(in_ch),
            nn.SiLU(inplace=True),
            nn.Conv2d(in_ch, out_ch, 1)  # 逐点卷积
        )

    def forward(self, x):
        return self.conv(x)


class SpatialAttention(nn.Module):
    def __init__(self,in_channels,kernel_size=7):
        super(SpatialAttention, self).__init__()
        assert kernel_size in (3, 7), 'kernel size must be 3 or 7'
        padding = 3 if kernel_size == 7 else 1

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=padding, bias=False)
        self.sigmoid = nn.Sigmoid()
        self.conv2 =nn.Sequential(nn.Conv2d(in_channels, in_channels//2, kernel_size, padding=padding, bias=False),
                                    nn.BatchNorm2d(in_channels//2),
                                   nn.ReLU()
                                  )

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        x_cat = torch.cat([avg_out, max_out], dim=1)
        spatial_weights = self.sigmoid(self.conv1(x_cat))  # 生成权重
        return self.conv2(x * spatial_weights)

class FPN_PAN_Enhanced(nn.Module):
    def __init__(self, channels=None, depth_multiple=0.33):
        super().__init__()
        if channels is None:
            channels = [256, 512, 1024]

        # 改进1：双线性插值 + 可学习上采样
        self.upsample_p5 = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            Conv(channels[2], channels[1], 3)  # 通道对齐
        )
        self.upsample_p4 = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            Conv(channels[1], channels[0], 3)
        )
        self.before_add = nn.Sequential(
            nn.Conv2d(channels[0],channels[0],3,2,1),
            nn.ReLU(),
            nn.Conv2d(channels[0], channels[0], 3, 2, 1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.LayerNorm(512),
        )

    def forward(self, x):
        p3, p4, p5 = x
        # 自顶向下路径（改进后）
        p5_up = self.upsample_p5(p5)
        p4 =torch.concat([p4, p5_up],dim=1)
        spatialattention1=SpatialAttention(2048)
        p4 = spatialattention1(p4)

        # 自底向上路径（改进后）
        p4_up = self.upsample_p4(p4)
        p3 = torch.concat([p3, p4_up],dim=1)
        spatialattention2 = SpatialAttention(1024)
        p3 = spatialattention2(p3)
        out = self.before_add(p3)
        # 分类输出（新增）
        return out


class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.encoder1 = ResNet50()
        self.encoder2 = ResNet50()
        self.fc_block = nn.Sequential(
            nn.Linear(42, 128),
            nn.BatchNorm1d(128),  # 替换 Dropout 为 BatchNorm1d
            nn.GELU(),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256),  # 替换 Dropout 为 BatchNorm1d
            nn.GELU(),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512)  # 替换 Dropout 为 BatchNorm1d
        )
        self.class_head = nn.Sequential(
            nn.Linear(1024, 512),
            nn.BatchNorm1d(512),  # 替换 Dropout 为 BatchNorm1d
            nn.GELU(),
            nn.Linear(512, 4)
        )
        # self.data_encoder = FCBlock(input_dim=42, hidden_dims=[128, 256], output_dim=512)

    def forward(self, x1,x2,x3):
        DIFF_tezheng1, DIFF_tezheng2, DIFF_tezheng3 = self.encoder1(x1)
        WNB_tezheng1, WNB_tezheng2, WNB_tezheng3 = self.encoder2(x2)
        merge_p3 = torch.cat([DIFF_tezheng1,WNB_tezheng1],dim=1)
        channelattention1=ChannelAttention(512)
        merge_p3 = channelattention1(merge_p3)

        merge_p4 = torch.cat([DIFF_tezheng2, WNB_tezheng2], dim=1)
        channelattention2= ChannelAttention(1024)
        merge_p4 = channelattention2(merge_p4)

        merge_p5 = torch.cat([DIFF_tezheng3, WNB_tezheng3], dim=1)
        channelattention3= ChannelAttention(2048)
        merge_p5 = channelattention3(merge_p5)
        fpn_pan=FPN_PAN_Enhanced([512,1024,2048])
        fpn_pan_out=fpn_pan([merge_p3,merge_p4,merge_p5])
        tabular_feat = self.fc_block(x3)

        # 最终分类

        mix = torch.cat([fpn_pan_out,tabular_feat],-1)

        #加入通道注意力机制
        # transformer = TransformerEncoder(d_model=1024).to(device)
        # out = transformer(out)
        out = self.class_head(mix)  # 特征融合
        return out




# 创建ResNet50模型
model = Model()

# 前向传播示例
input_tensor1 = torch.randn(2, 3, 256, 256)
input_tensor2 = torch.randn(2, 3, 256, 256)
input_tensor3 = torch.randn(2, 42)
output = model(input_tensor1,input_tensor2,input_tensor3)  # shape: [1, 10]

# 查看参数量
print(sum(p.numel() for p in model.parameters())/1024)  # 约25.6M